import time
import sys, pygame

size = 500.0
c = 10.0
h = 1.0
time = (h / c) / 1.1
w_width = size
w_height = 400.0
max_height = 1.0

def update(velocities, heights, size, time):
    new = []
    for i in range(0, int(size)):
        left = heights[i]
        if i > 0:
            left = heights[i - 1]
        right = heights[i]
        if i < size - 1:
            right = heights[i + 1]
        f = c * c * (left + right - 2 * heights[i]) / (h * h)
        velocities[i] += f * time
        new.append(heights[i] + velocities[i] * time)
    for i in range(0, int(size)):
        heights[i] = new[i]

def print_waves(heights, size, screen):
    width = w_width / size
    for i in range(0, int(size)):
        if width >= 1.0:
            rect = pygame.Rect(i * width, w_height - heights[i] / max_height * w_height,
                           width, heights[i] / max_height * w_height)
            surface = pygame.Surface((width, heights[i] / max_height * w_height))
        else:
            rect = pygame.Rect(i * width, w_height - heights[i] / max_height * w_height,
                               1, heights[i] / max_height * w_height)
            surface = pygame.Surface((1, heights[i] / max_height * w_height))
        surface.fill(pygame.Color(0, 0, 255, 0))
        screen.blit(surface, rect)

velocities = []
heights = []

for i in range(0, int(size)):
    velocities.append(0.0)
    heights.append(0.3)

#for i in range(0, int(size) / 10):
#    heights[i] = 0.35

screen = pygame.display.set_mode((int(w_width), int(w_height)))

while 1:
    screen.fill(pygame.Color(0, 0, 0, 0))
    for event in pygame.event.get():
        if event.type == pygame.QUIT: sys.exit()
    update(velocities, heights, size, time)
    print_waves(heights, size, screen)
    pygame.display.flip()
